package com.ai.controller;

import cn.hutool.core.util.StrUtil;
import com.ai.chat.RedisChatMemory;
import com.ai.service.ChatSessionService;
import com.ai.entity.ChatSession;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.UUID;

@RestController
@RequestMapping("ai/chat")
@RequiredArgsConstructor
public class AIChatController {

    private final ChatClient chatClient;
    private final RedisChatMemory redisChatMemory;
    private final ChatSessionService chatSessionService;
    String userId = "1";

    /**
     * 流式回答聊天（会话）
     */
    @GetMapping(value = "/generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatClientResponse> generateStream(@RequestParam(value = "message") String message, @RequestParam(required = false) String sessionId) {
        // 默认生成一个会话
        String conversationId = getSessionId(message, sessionId);
        return chatClient.prompt().user(message).advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId)).stream().chatClientResponse();
    }

    /**
     * 当前会话的所有信息
     *
     * @param sessionId 会话ID
     */
    @GetMapping("/messages")
    public ResponseEntity<List<Message>> getMessages(@RequestParam String sessionId) {
        return ResponseEntity.ok(redisChatMemory.get(sessionId));
    }

    /**
     * 当前登录人的会话列表
     */
    @GetMapping("/sessions")
    public ResponseEntity<List<ChatSession>> getSessions() {
        return ResponseEntity.ok(chatSessionService.getSessions(userId));
    }

    /**
     * 普通聊天（会话）
     */
    @GetMapping(value = "/generate")
    public ResponseEntity<String> generate(@RequestParam(value = "message") String message, @RequestParam String sessionId) {
        String conversationId = getSessionId(message, sessionId);
        return ResponseEntity.ok(chatClient.prompt().user(message).advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId)).call().content());
    }

    private String getSessionId(@RequestParam("message") String message, @RequestParam(required = false) String sessionId) {
        if (StrUtil.isBlank(sessionId)) {
            sessionId = UUID.randomUUID().toString();

            ChatSession chatSession = new ChatSession().setSessionId(sessionId).setSessionName(message.length() >= 15 ? message.substring(0, 15) : message);
            chatSessionService.saveSession(chatSession, userId);
        }
        return sessionId;
    }
}
